from functools import wraps
from fastapi import HTTPException


def return_response():
    """
    装饰器：统一返回格式。
    返回格式：
    {
        "code": 200,             # 状态码，200表示成功，非200表示失败
        "message": "success",    # 描述信息
        "total": 100,            # 总记录数
        "page": 1,               # 当前页码
        "size": 10,              # 每页数据条数
        "data": [...]            # 返回的具体数据
    }
    """

    def response_decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            try:
                # 执行被装饰的函数
                data = await func(*args, **kwargs)

                # 判断返回的数据是否包含分页信息
                if isinstance(data, dict) and 'total' in data and 'page' in data and 'size' in data:
                    # 返回包含分页信息的标准格式
                    return {
                        "code": 200,
                        "message": "success",
                        "total": data.get('total', 0),
                        "page": data.get('page', 1),
                        "size": data.get('size', 10),
                        "data": data.get('data', [])
                    }

                # 如果没有分页信息，则直接返回标准格式
                return {
                    "code": 200,
                    "message": "success",
                    "data": data
                }

            except HTTPException as e:
                # 捕获 FastAPI 抛出的 HTTP 异常
                return {
                    "code": e.status_code,
                    "message": e.detail,
                    "data": None
                }
            except Exception as e:
                # 捕获其他异常
                return {
                    "code": 500,
                    "message": f"Internal Server Error: {str(e)}",
                    "data": None
                }

        return wrapper

    return response_decorator
